"""
ChromaDB 数据库查看器
提供友好的 Web 界面查看数据库内容
"""
import sys
import json
import chromadb
from datetime import datetime
from flask import Flask, jsonify, render_template_string

app = Flask(__name__)

# 初始化 ChromaDB
CHROMA_PERSIST_DIR = "./chroma"
COLLECTION_NAME = "wind_blade_defects"

try:
    client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR)
    collection = client.get_collection(name=COLLECTION_NAME)
    print(f"✅ 成功连接到 ChromaDB 集合: {COLLECTION_NAME}")
except Exception as e:
    print(f"❌ 连接失败: {e}")
    collection = None

HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>ChromaDB 数据库查看器</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Microsoft YaHei', sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }
        
        .container {
            max-width: 1400px;
            margin: 0 auto;
        }
        
        .header {
            background: white;
            border-radius: 16px;
            padding: 30px;
            margin-bottom: 24px;
            box-shadow: 0 10px 40px rgba(0, 0, 0, 0.1);
        }
        
        .header h1 {
            font-size: 32px;
            color: #1f2937;
            margin-bottom: 12px;
            display: flex;
            align-items: center;
            gap: 12px;
        }
        
        .header .subtitle {
            color: #6b7280;
            font-size: 16px;
        }
        
        .stats {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
            gap: 16px;
            margin-bottom: 24px;
        }
        
        .stat-card {
            background: white;
            border-radius: 12px;
            padding: 24px;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
            transition: transform 0.2s;
        }
        
        .stat-card:hover {
            transform: translateY(-4px);
        }
        
        .stat-icon {
            font-size: 32px;
            margin-bottom: 12px;
        }
        
        .stat-value {
            font-size: 36px;
            font-weight: 700;
            color: #1f2937;
            margin-bottom: 8px;
        }
        
        .stat-label {
            color: #6b7280;
            font-size: 14px;
            text-transform: uppercase;
            letter-spacing: 0.5px;
        }
        
        .controls {
            background: white;
            border-radius: 12px;
            padding: 20px;
            margin-bottom: 24px;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
            display: flex;
            gap: 12px;
            flex-wrap: wrap;
            align-items: center;
        }
        
        .controls input {
            flex: 1;
            min-width: 250px;
            padding: 12px 16px;
            border: 2px solid #e5e7eb;
            border-radius: 8px;
            font-size: 14px;
        }
        
        .controls input:focus {
            outline: none;
            border-color: #667eea;
        }
        
        .btn {
            padding: 12px 24px;
            border: none;
            border-radius: 8px;
            font-size: 14px;
            font-weight: 600;
            cursor: pointer;
            transition: all 0.2s;
            display: inline-flex;
            align-items: center;
            gap: 8px;
        }
        
        .btn-primary {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
        }
        
        .btn-primary:hover {
            transform: translateY(-2px);
            box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4);
        }
        
        .btn-secondary {
            background: #f3f4f6;
            color: #374151;
        }
        
        .btn-secondary:hover {
            background: #e5e7eb;
        }
        
        .conversations {
            display: grid;
            gap: 16px;
        }
        
        .conversation-card {
            background: white;
            border-radius: 12px;
            padding: 24px;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
            transition: all 0.3s;
            border-left: 4px solid #667eea;
        }
        
        .conversation-card:hover {
            box-shadow: 0 8px 24px rgba(0, 0, 0, 0.15);
            transform: translateX(4px);
        }
        
        .conversation-header {
            display: flex;
            justify-content: space-between;
            align-items: flex-start;
            margin-bottom: 16px;
            padding-bottom: 16px;
            border-bottom: 2px solid #f3f4f6;
        }
        
        .conversation-id {
            font-size: 12px;
            color: #9ca3af;
            font-family: 'Courier New', monospace;
        }
        
        .conversation-time {
            font-size: 13px;
            color: #6b7280;
            display: flex;
            align-items: center;
            gap: 6px;
        }
        
        .message-section {
            margin-bottom: 16px;
        }
        
        .message-label {
            font-size: 12px;
            font-weight: 700;
            color: #374151;
            text-transform: uppercase;
            letter-spacing: 0.5px;
            margin-bottom: 8px;
            display: flex;
            align-items: center;
            gap: 6px;
        }
        
        .message-content {
            background: #f9fafb;
            padding: 16px;
            border-radius: 8px;
            color: #1f2937;
            line-height: 1.6;
            white-space: pre-wrap;
            word-break: break-word;
        }
        
        .user-message {
            border-left: 3px solid #3b82f6;
        }
        
        .bot-message {
            border-left: 3px solid #10b981;
        }
        
        .metadata {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(150px, 1fr));
            gap: 12px;
            margin-top: 16px;
            padding-top: 16px;
            border-top: 2px solid #f3f4f6;
        }
        
        .metadata-item {
            font-size: 13px;
        }
        
        .metadata-label {
            color: #6b7280;
            font-weight: 600;
            margin-bottom: 4px;
        }
        
        .metadata-value {
            color: #374151;
            font-family: 'Courier New', monospace;
        }
        
        .badge {
            display: inline-block;
            padding: 4px 12px;
            border-radius: 12px;
            font-size: 11px;
            font-weight: 600;
            text-transform: uppercase;
        }
        
        .badge-success {
            background: #d1fae5;
            color: #065f46;
        }
        
        .badge-info {
            background: #dbeafe;
            color: #1e40af;
        }
        
        .loading {
            text-align: center;
            padding: 60px;
            color: white;
            font-size: 18px;
        }
        
        .empty-state {
            background: white;
            border-radius: 12px;
            padding: 60px;
            text-align: center;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
        }
        
        .empty-icon {
            font-size: 64px;
            margin-bottom: 16px;
            opacity: 0.6;
        }
        
        .empty-title {
            font-size: 20px;
            font-weight: 700;
            color: #374151;
            margin-bottom: 8px;
        }
        
        .empty-text {
            color: #6b7280;
            font-size: 14px;
        }
        
        .vector-preview {
            margin-top: 12px;
            padding: 12px;
            background: #f3f4f6;
            border-radius: 6px;
            font-size: 11px;
            font-family: 'Courier New', monospace;
            color: #6b7280;
            max-height: 100px;
            overflow: auto;
        }
        
        @media (max-width: 768px) {
            .stats {
                grid-template-columns: 1fr;
            }
            
            .controls {
                flex-direction: column;
            }
            
            .controls input {
                width: 100%;
            }
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>
                <span>🗄️</span>
                <span>ChromaDB 数据库查看器</span>
            </h1>
            <p class="subtitle">实时查看向量数据库中的聊天历史记录</p>
        </div>
        
        <div class="stats" id="stats"></div>
        
        <div class="controls">
            <input type="text" id="searchInput" placeholder="🔍 搜索对话内容...">
            <button class="btn btn-primary" onclick="loadData()">
                <span>🔄</span>
                <span>刷新数据</span>
            </button>
            <button class="btn btn-secondary" onclick="exportData()">
                <span>📥</span>
                <span>导出 JSON</span>
            </button>
        </div>
        
        <div id="content" class="loading">
            <div>⏳ 正在加载数据...</div>
        </div>
    </div>
    
    <script>
        let allData = [];
        
        async function loadData() {
            try {
                document.getElementById('content').innerHTML = '<div class="loading">⏳ 正在加载数据...</div>';
                
                const response = await fetch('/api/data');
                const data = await response.json();
                
                allData = data.conversations;
                updateStats(data.stats);
                displayConversations(allData);
            } catch (error) {
                document.getElementById('content').innerHTML = `
                    <div class="empty-state">
                        <div class="empty-icon">❌</div>
                        <div class="empty-title">加载失败</div>
                        <div class="empty-text">${error.message}</div>
                    </div>
                `;
            }
        }
        
        function updateStats(stats) {
            const statsHTML = `
                <div class="stat-card">
                    <div class="stat-icon">📊</div>
                    <div class="stat-value">${stats.total}</div>
                    <div class="stat-label">总对话数</div>
                </div>
                <div class="stat-card">
                    <div class="stat-icon">📅</div>
                    <div class="stat-value">${stats.today}</div>
                    <div class="stat-label">今日对话</div>
                </div>
                <div class="stat-card">
                    <div class="stat-icon">📈</div>
                    <div class="stat-value">${stats.week}</div>
                    <div class="stat-label">本周对话</div>
                </div>
                <div class="stat-card">
                    <div class="stat-icon">💾</div>
                    <div class="stat-value">${(stats.dbSize / 1024).toFixed(1)} KB</div>
                    <div class="stat-label">数据库大小</div>
                </div>
            `;
            document.getElementById('stats').innerHTML = statsHTML;
        }
        
        function displayConversations(conversations) {
            if (conversations.length === 0) {
                document.getElementById('content').innerHTML = `
                    <div class="empty-state">
                        <div class="empty-icon">💬</div>
                        <div class="empty-title">暂无对话记录</div>
                        <div class="empty-text">开始使用聊天功能后，对话记录将显示在这里</div>
                    </div>
                `;
                return;
            }
            
            const html = conversations.map(conv => {
                const timestamp = conv.metadata.timestamp ? new Date(conv.metadata.timestamp).toLocaleString('zh-CN') : '未知时间';
                const hasResults = conv.metadata.hasResults ? '有图片' : '无图片';
                const provider = conv.metadata.provider || '未知';
                
                return `
                    <div class="conversation-card">
                        <div class="conversation-header">
                            <div class="conversation-id">ID: ${conv.id.substring(0, 8)}...</div>
                            <div class="conversation-time">
                                <span>🕒</span>
                                <span>${timestamp}</span>
                            </div>
                        </div>
                        
                        <div class="message-section">
                            <div class="message-label">
                                <span>👤</span>
                                <span>用户问题</span>
                            </div>
                            <div class="message-content user-message">
                                ${conv.metadata.userMessage || '无内容'}
                            </div>
                        </div>
                        
                        <div class="message-section">
                            <div class="message-label">
                                <span>🤖</span>
                                <span>AI 回复</span>
                            </div>
                            <div class="message-content bot-message">
                                ${conv.metadata.botResponse || '无内容'}
                            </div>
                        </div>
                        
                        <div class="metadata">
                            <div class="metadata-item">
                                <div class="metadata-label">模型提供者</div>
                                <div class="metadata-value">
                                    <span class="badge badge-info">${provider}</span>
                                </div>
                            </div>
                            <div class="metadata-item">
                                <div class="metadata-label">检索结果</div>
                                <div class="metadata-value">
                                    <span class="badge ${hasResults === '有图片' ? 'badge-success' : 'badge-info'}">${hasResults}</span>
                                </div>
                            </div>
                            <div class="metadata-item">
                                <div class="metadata-label">向量维度</div>
                                <div class="metadata-value">${conv.embedding ? conv.embedding.length : 'N/A'}</div>
                            </div>
                        </div>
                        
                        ${conv.embedding ? `
                            <div class="vector-preview">
                                <strong>向量预览:</strong> [${conv.embedding.slice(0, 10).map(v => v.toFixed(4)).join(', ')}...]
                            </div>
                        ` : ''}
                    </div>
                `;
            }).join('');
            
            document.getElementById('content').innerHTML = `<div class="conversations">${html}</div>`;
        }
        
        function exportData() {
            const dataStr = JSON.stringify(allData, null, 2);
            const dataBlob = new Blob([dataStr], { type: 'application/json' });
            const url = URL.createObjectURL(dataBlob);
            const link = document.createElement('a');
            link.href = url;
            link.download = `chromadb_export_${new Date().getTime()}.json`;
            link.click();
            URL.revokeObjectURL(url);
        }
        
        // 搜索功能
        document.getElementById('searchInput').addEventListener('input', (e) => {
            const query = e.target.value.toLowerCase();
            if (!query) {
                displayConversations(allData);
                return;
            }
            
            const filtered = allData.filter(conv => 
                conv.metadata.userMessage?.toLowerCase().includes(query) ||
                conv.metadata.botResponse?.toLowerCase().includes(query) ||
                conv.id.includes(query)
            );
            displayConversations(filtered);
        });
        
        // 初始加载
        loadData();
        
        // 自动刷新（每30秒）
        setInterval(loadData, 30000);
    </script>
</body>
</html>
"""

@app.route('/')
def index():
    return render_template_string(HTML_TEMPLATE)

@app.route('/api/data')
def get_data():
    if not collection:
        return jsonify({
            'error': 'ChromaDB 未连接',
            'conversations': [],
            'stats': {
                'total': 0,
                'today': 0,
                'week': 0,
                'dbSize': 0
            }
        }), 500
    
    try:
        # 获取所有数据
        results = collection.get(
            include=['embeddings', 'metadatas', 'documents']
        )
        
        conversations = []
        for i in range(len(results['ids'])):
            conversations.append({
                'id': results['ids'][i],
                'embedding': results['embeddings'][i] if results['embeddings'] else None,
                'metadata': results['metadatas'][i] if results['metadatas'] else {},
                'document': results['documents'][i] if results['documents'] else ''
            })
        
        # 按时间倒序排列
        conversations.sort(
            key=lambda x: x['metadata'].get('timestamp', ''),
            reverse=True
        )
        
        # 计算统计信息
        now = datetime.now()
        today_count = sum(1 for c in conversations if c['metadata'].get('timestamp', '').startswith(now.strftime('%Y-%m-%d')))
        
        week_ago = datetime.now().timestamp() - (7 * 24 * 60 * 60)
        week_count = sum(1 for c in conversations if datetime.fromisoformat(c['metadata'].get('timestamp', '1970-01-01')).timestamp() > week_ago)
        
        # 获取数据库文件大小
        import os
        db_size = 0
        try:
            db_path = f"{CHROMA_PERSIST_DIR}/chroma.sqlite3"
            if os.path.exists(db_path):
                db_size = os.path.getsize(db_path)
        except:
            pass
        
        return jsonify({
            'conversations': conversations,
            'stats': {
                'total': len(conversations),
                'today': today_count,
                'week': week_count,
                'dbSize': db_size
            }
        })
    except Exception as e:
        return jsonify({
            'error': str(e),
            'conversations': [],
            'stats': {
                'total': 0,
                'today': 0,
                'week': 0,
                'dbSize': 0
            }
        }), 500

if __name__ == '__main__':
    print('=' * 60)
    print('🗄️  ChromaDB 数据库查看器')
    print('=' * 60)
    print(f'📁 数据库位置: {CHROMA_PERSIST_DIR}')
    print(f'📊 集合名称: {COLLECTION_NAME}')
    print('')
    print('🌐 访问地址: http://localhost:5555')
    print('')
    print('💡 提示:')
    print('   - 打开浏览器访问上述地址查看数据库内容')
    print('   - 支持实时搜索和数据导出')
    print('   - 每30秒自动刷新数据')
    print('=' * 60)
    print('')
    
    app.run(host='0.0.0.0', port=5555, debug=False)
